﻿using System;
using Mono.Cecil.Cil;
using Telerik.JustDecompiler.Cil;
using System.Collections.Generic;
using Mono.Cecil;

namespace Telerik.JustDecompiler.Decompiler.StateMachines
{
    class AsyncStateControllerRemover : StateControllerRemover
    {
        private readonly VariableReference doFinallyVariable;
        private readonly AsyncStateMachineVersion version;

        private InstructionBlock debugStateCheckBlock;
        private int stateCheckOffset;

        public bool FoundControllerBlocks
        {
            get
            {
                return this.debugStateCheckBlock != null || this.BlocksMarkedForRemoval.Count > firstControllerBlock;
            }
        }

        public AsyncStateControllerRemover(MethodSpecificContext methodContext, FieldDefinition stateField, VariableReference doFinallyVariable, AsyncStateMachineVersion version)
            :base(methodContext, stateField)
        {
            this.doFinallyVariable = doFinallyVariable;
            this.version = version;
        }

        public override bool RemoveStateMachineController()
        {
            if (version == AsyncStateMachineVersion.V1)
            {
                if (IsDebugCheckStateBlock(theCFG.Blocks[0]))
                {
                    if (!GetStateFieldAndVariable())
                    {
                        return false;
                    }
                    SkipFirstBlock();
                    return RemoveControllerChain();
                }
                else if (IsDoFinallySetBlock(theCFG.Blocks[0]))
                {
                    SkipFirstBlock();
                }

                return base.RemoveStateMachineController();
            }
            else
            {
                if (ContainsStateFieldLoad(theCFG.Blocks[0]))
                {
                    if (!GetStateFieldAndVariable())
                    {
                        return false;
                    }

                    SkipFirstBlock();
                }

                return RemoveControllerChain();
            }
        }

        private void SkipFirstBlock()
        {
            this.BlocksMarkedForRemoval.Add(theCFG.Blocks[0]);
            firstControllerBlock = 1;
        }

        protected override bool IsUnconditionalBranchBlock(InstructionBlock theBlock)
        {
            return base.IsUnconditionalBranchBlock(theBlock) || IsDummyStateControllerBlock(theBlock) || IsNopBlock(theBlock);
        }

        /// <summary>
        /// Checks whether the specified block is a dummy controller block.
        /// </summary>
        /// <remarks>
        /// This block looks like a state controller block with the difference that it ends with two pop instructions.
        /// Generated by the C# compiler in release mode.
        /// Update: With C# 6.0 the structure of the dummy state controller blocks has changed. The new one is load
        /// of the state variable, followed by pop and nop instructions.
        /// </remarks>
        /// <param name="theBlock"></param>
        /// <returns></returns>
        private bool IsDummyStateControllerBlock(InstructionBlock theBlock)
        {
            if (this.version == AsyncStateMachineVersion.V1)
            {
                if (!ContainsStateFieldLoad(theBlock))
                {
                    return false;
                }

                Instruction currentInstruction = theBlock.Last;
                for (int i = 0; i < 2; i++)
                {
                    if (currentInstruction == theBlock.First || currentInstruction.OpCode.Code != Code.Pop)
                    {
                        return false;
                    }
                    currentInstruction = currentInstruction.Previous;
                }

                toBeRemoved.Add(theBlock);
                return true;
            }
            else
            {
                Instruction current = theBlock.First;
                VariableReference variable;
                if (!TryGetVariableFromInstruction(current, out variable) || variable != stateVariable)
                {
                    return false;
                }

                if (current.Next.OpCode.Code != Code.Pop || current.Next.Next.OpCode.Code != Code.Nop)
                {
                    return false;
                }

                if (current.Next.Next != theBlock.Last)
                {
                    return false;
                }

                return true;
            }
        }

        /// <summary>
        /// Checks whether the specified block is a nop block.
        /// </summary>
        /// <param name="theBlock"></param>
        /// <returns></returns>
        private bool IsNopBlock(InstructionBlock theBlock)
        {
            return IsNopTillEnd(theBlock, theBlock.First);
        }
        
        /// <summary>
        /// Checks whether the specified block contains only assignment of the doFinallyBodies variable.
        /// </summary>
        /// <param name="theBlock"></param>
        /// <returns></returns>
        private bool IsDoFinallySetBlock(InstructionBlock theBlock)
        {
            if (!BeginsWithDoFinallySet(theBlock))
            {
                return false;
            }

            Instruction currentInstruction = theBlock.First.Next;
            return currentInstruction == theBlock.Last || IsNopTillEnd(theBlock, currentInstruction.Next);
        }

        /// <summary>
        /// Checks whether the specified block contains only nop instructions starting from the specified instruction until the block's end.
        /// </summary>
        /// <param name="theBlock"></param>
        /// <param name="currentInstruction"></param>
        /// <returns></returns>
        private bool IsNopTillEnd(InstructionBlock theBlock, Instruction currentInstruction)
        {
            while (currentInstruction != theBlock.Last)
            {
                if (currentInstruction.OpCode.Code != Code.Nop)
                {
                    return false;
                }
                currentInstruction = currentInstruction.Next;
            }

            return currentInstruction.OpCode.Code == Code.Nop;
        }

        /// <summary>
        /// Checks whether the specified block is a state controller block generated by the C# compiler in debug mode.
        /// </summary>
        /// <remarks>
        /// Pattern:
        /// ldc.i4.1
        /// stloc* doFinallyVariable
        /// ldarg.0
        /// ldfld stateField
        /// ...
        /// ldc.i4.s -3
        /// (sub) - missing if next is beq*
        /// switch || beq*
        /// 
        /// There is only one such block generated at the begining of the MoveNext method.
        /// </remarks>
        /// <param name="theBlock"></param>
        /// <returns></returns>
        private bool IsDebugCheckStateBlock(InstructionBlock theBlock)
        {
            if (!BeginsWithDoFinallySet(theBlock) || !ContainsStateFieldLoad(theBlock))
            {
                return false;
            }

            Instruction currentInstruction = theBlock.Last;
            if (!this.IsBeqInstruction(currentInstruction))
            {
                if (currentInstruction.OpCode.Code != Code.Switch)
                {
                    return false;
                }

                currentInstruction = currentInstruction.Previous;
                if (currentInstruction.OpCode.Code != Code.Sub)
                {
                    return false;
                }
            }

            currentInstruction = currentInstruction.Previous;
            int operand;
            if (!StateMachineUtilities.TryGetOperandOfLdc(currentInstruction, out operand) || operand > -3)
            {
                return false;
            }

            debugStateCheckBlock = theBlock;
            stateCheckOffset = -operand;
            return true;
        }

        /// <summary>
        /// Checks whether the specified block begins with setting the doFinallyBodies variable to true.
        /// </summary>
        /// <param name="theBlock"></param>
        /// <returns></returns>
        private bool BeginsWithDoFinallySet(InstructionBlock theBlock)
        {
            Instruction currentInstruction = theBlock.First;
            if (currentInstruction == theBlock.Last || currentInstruction.OpCode.Code != Code.Ldc_I4_1)
            {
                return false;
            }

            VariableReference foundVariable;
            return StateMachineUtilities.TryGetVariableFromInstruction(currentInstruction.Next, methodContext.Body.Variables, out foundVariable) &&
                foundVariable == doFinallyVariable;
        }

        /// <summary>
        /// Initializes the queue that is used for traversing the state controller blocks.
        /// </summary>
        /// <remarks>
        /// If the body of the MoveNext method starts with a debugCheckStateBlock then we add to the queue the beginning of each non-negative state
        /// (the states that are numbered with a non-negative integer).
        /// </remarks>
        /// <returns></returns>
        protected override Queue<InstructionBlock> InitializeTheTraversalQueue()
        {
            Queue<InstructionBlock> theQueue = base.InitializeTheTraversalQueue();

            SwitchData switchData;
            if (debugStateCheckBlock == null || !theCFG.SwitchBlocksInformation.TryGetValue(debugStateCheckBlock, out switchData))
            {
                return theQueue;
            }

            for (int i = 0; i + stateCheckOffset < switchData.OrderedCasesArray.Length; i++)
            {
                stateToStartBlock[i] = SkipBranchChain(switchData.OrderedCasesArray[i + stateCheckOffset]);
                theQueue.Enqueue(stateToStartBlock[i]);
            }

            return theQueue;
        }
    }
}
